from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.serialization import Encoding
from cryptography.hazmat.primitives.asymmetric import padding
import base64
from rsa import PublicKey, common, transform, core
# from rsa.types import RSAOAEPHashAlgorithm

# 生成一个2048位的私钥和公钥
private_key = rsa.generate_private_key(
    public_exponent=65537,
    key_size=2048,
    backend=default_backend()
)

public_key = private_key.public_key()

# 将私钥和公钥序列化为PEM格式
private_pem = private_key.private_bytes(
    encoding=Encoding.PEM,
    format=serialization.PrivateFormat.PKCS8,
    encryption_algorithm=serialization.NoEncryption()
)

public_pem = public_key.public_bytes(
    encoding=Encoding.PEM,
    format=serialization.PublicFormat.PKCS1
)

# 加密和解密的示例
message = b"secret data"

# 对消息进行加密
ciphertext = public_key.encrypt(
    message,
    padding.OAEP(
        mgf=padding.MGF1(algorithm=rsa.types.RSAOAEPHashAlgorithm.SHA256),
        algorithm=rsa.types.RSAOAEPHashAlgorithm.SHA256,
        label=None
    )
)

# 对密文进行解密
plaintext = private_key.decrypt(
    ciphertext,
    padding.OAEP(
        mgf=padding.MGF1(algorithm=rsa.types.RSAOAEPHashAlgorithm.SHA256),
        algorithm=rsa.types.RSAOAEPHashAlgorithm.SHA256,
        label=None
    )
)

# 打印结果
print("Original message:", message)
print("Encrypted message (Base64):", base64.b64encode(ciphertext).decode('utf-8'))
print("Decrypted message:", plaintext)